from collections.abc import Generator
from typing import Any
from dify_plugin import Tool
from dify_plugin.entities.tool import ToolInvokeMessage

import requests
import csv
import os



class YearreportcrawlerTool(Tool):
    def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage]:

        company_name = tool_parameters.get('company_name')
        # stock_code = tool_parameters.get('stock_code')
        stock_code = get_stock_code(company_name)
        if stock_code:
            print(f"成功获取 {company_name} 的股票代码: {stock_code}")
            print(f"Fetching annual reports for {company_name} ({stock_code})...")
        else:
            print(f"未找到 {company_name} 的股票代码。")

        start_date = tool_parameters.get('start_date')
        end_date = tool_parameters.get('end_date')

        # output_folder = tool_parameters.get('output_folder')
        output_folder = "annual_reports"
        pdf_folder = os.path.join(output_folder, "pdf_files")
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        if not os.path.exists(pdf_folder):
            os.makedirs(pdf_folder)
        reports = get_annual_reports(stock_code, start_date, end_date)
        if reports:
            filename = os.path.join(output_folder, f"{company_name}_annual_reports_{start_date[:4]}-{end_date[:4]}.csv")
            save_to_csv(reports, filename)

            # 下载PDF文件
            for item in reports:
                pdf_url = f"http://static.cninfo.com.cn/{item['adjunctUrl']}"
                # 生成更友好的文件名
                pdf_filename = f"{item['secCode']}_{item['announcementTitle'][:30]}.pdf"
                pdf_save_path = os.path.join(pdf_folder, pdf_filename)
                download_pdf(pdf_url, pdf_save_path)
        else:
            print(f"No data found for {company_name} ({stock_code}). Skipping...")
        yield self.create_json_message({
            "result": reports
        })

def get_stock_code(company_name):
    url = "http://www.cninfo.com.cn/new/hisAnnouncement/query"
    headers = {
        "Accept": "application/json, text/javascript, */*; q=0.01",
        "Content-Type": "application/x-www-form-urlencoded; charset=UTF-8",
        "Origin": "http://www.cninfo.com.cn",
        "Referer": "http://www.cninfo.com.cn/new/commonUrl/pageOfSearch?url=disclosure/list/search",
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36",
        "X-Requested-With": "XMLHttpRequest"
    }
    data = {
        "pageNum": 1,
        "pageSize": 1,
        "column": "szse",
        "tabName": "fulltext",
        "plate": "",
        "stock": "",
        "searchkey": company_name,
        "secid": "",
        "category": "category_ndbg_szsh",
        "trade": "",
        "seDate": "2020-01-01~2024-12-31",
        "sortName": "",
        "sortType": "",
        "isHLtitle": "false"
    }
    response = requests.post(url, headers=headers, data=data)
    if response.status_code == 200:
        response_data = response.json()
        print(response_data)
        if "announcements" in response_data and len(response_data["announcements"]) > 0:
            sec_code = response_data["announcements"][0]["secCode"]
            stock_suffix = response_data["announcements"][0]["orgId"]
            stock_code = f"{sec_code},{stock_suffix}"
            return stock_code
    return None

def get_annual_reports(stock_code, start_date, end_date):
    url = "http://www.cninfo.com.cn/new/hisAnnouncement/query"
    headers = {
        "Accept": "application/json, text/javascript, */*; q=0.01",
        "Content-Type": "application/x-www-form-urlencoded; charset=UTF-8",
        "Origin": "http://www.cninfo.com.cn",
        "Referer": "http://www.cninfo.com.cn/new/commonUrl/pageOfSearch?url=disclosure/list/search",
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36",
        "X-Requested-With": "XMLHttpRequest"
    }
    data = {
        "pageNum": 1,
        "pageSize": 30,
        "column": "szse",
        "tabName": "fulltext",
        "plate": "",
        "stock": stock_code,
        "searchkey": "",
        "secid": "",
        "category": "category_ndbg_szsh",
        "trade": "",
        "seDate": f"{start_date}~{end_date}",
        "sortName": "",
        "sortType": "",
        "isHLtitle": "false"
    }
    response = requests.post(url, headers=headers, data=data)
    if response.status_code == 200:
        response_data = response.json()
        if "announcements" in response_data:
            return response_data["announcements"]
        else:
            print(f"No announcements found for {stock_code}. Response: {response_data}")
    else:
        print(f"Failed to fetch data for {stock_code}. Status code: {response.status_code}")
    return None

# 定义保存到CSV的函数
def save_to_csv(data, filename):
    if data is None:
        print(f"No data to save. Skipping CSV creation for {filename}")
        return
    with open(filename, mode="w", newline="", encoding="utf-8") as file:
        writer = csv.writer(file)
        writer.writerow(["股票代码", "公司名称", "公告标题", "公告链接", "公告日期"])
        for item in data:
            writer.writerow([
                item["secCode"],
                item["secName"],
                item["announcementTitle"],
                f"http://static.cninfo.com.cn/{item['adjunctUrl']}",
                item["announcementTime"]
            ])
    print(f"Data saved to {filename}")

# 定义下载PDF文件的函数
def download_pdf(pdf_url, save_path):
    try:
        response = requests.get(pdf_url, stream=True)
        if response.status_code == 200:
            with open(save_path, "wb") as file:
                for chunk in response.iter_content(chunk_size=1024):
                    if chunk:
                        file.write(chunk)
            print(f"Downloaded {pdf_url} to {save_path}")
        else:
            print(f"Failed to download {pdf_url}. Status code: {response.status_code}")
    except Exception as e:
        print(f"Error downloading {pdf_url}: {e}")

# 主程序
if __name__ == "__main__":
    # 输入参数（可以改成动态输入）
    company_name = "永安林业"
    #stock_code = "601996,9900021478"  # 丰林集团的股票代码
    # stock_code = "603816,9900027317" #顾家家居
    # stock_code = "000663,gssz0000663"
    stock_code = get_stock_code(company_name)
    start_date = "2020-01-01"
    end_date = "2024-12-31"
    output_folder = "annual_reports"
    pdf_folder = os.path.join(output_folder, "pdf_files")

    # 创建输出文件夹
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    if not os.path.exists(pdf_folder):
        os.makedirs(pdf_folder)

    print(f"Fetching annual reports for {company_name} ({stock_code})...")
    reports = get_annual_reports(stock_code, start_date, end_date)
    if reports:
        filename = os.path.join(output_folder, f"{company_name}_annual_reports_{start_date[:4]}-{end_date[:4]}.csv")
        save_to_csv(reports, filename)

        # 下载PDF文件
        for item in reports:
            pdf_url = f"http://static.cninfo.com.cn/{item['adjunctUrl']}"
            # 生成更友好的文件名
            pdf_filename = f"{item['secCode']}_{item['announcementTitle'][:30]}.pdf"
            pdf_save_path = os.path.join(pdf_folder, pdf_filename)
            download_pdf(pdf_url, pdf_save_path)
    else:
        print(f"No data found for {company_name} ({stock_code}). Skipping...")